import numpy as np

from ..model import Vehicle
from ..sensor import RangeBearingSensor


class ExtendedKalmanFilter:
    def __init__(self, model: Vehicle, sensor: RangeBearingSensor, P_est, V_est, W_est) -> None:
        super().__init__()

        self._model = model
        self._sensor = sensor

        self._P_est = P_est
        self._V_est = V_est  # prediction noise
        self._W_est = W_est  # measure noise

    def step(self, v, z_measure):
        x_est = self._model.x
        x_pred = self._model.handle_state(self._model.F(x_est, v))
        Fx = self._model.Fx(x_est, v)
        Fv = self._model.Fv(x_est, v)
        P_pred = Fx @ self._P_est @ Fx.T + Fv @ self._V_est @ Fv.T

        y = self._sensor.handle_measure(z_measure - self._sensor.H(x_pred))
        Hx = self._sensor.Hx(x_pred)
        Hw = self._sensor.Hw(x_pred)
        S = Hx @ P_pred @ Hx.T + Hw @ self._W_est @ Hw.T
        K = P_pred @ Hx.T @ np.linalg.inv(S)

        x_est = self._model.handle_state(x_pred + K @ y)

        self._P_est[:, :] = P_pred - K @ Hx @ P_pred

        self._model.x = x_est
